from typing import Type
from pydantic import BaseModel, Field
from langchain.tools import BaseTool
from langchain.agents import AgentType, initialize_agent
from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
# 新版正确导入方式
from langchain_ollama import OllamaLLM # 用于基础LLM
from langchain_ollama  import ChatOllama  # 用于聊天模型

# 定义输入模型
class WebSearchInput(BaseModel):
    query: str = Field(..., description="搜索查询内容")

# 自定义网络搜索工具（修复字段覆盖问题）
class WebSearchTool(BaseTool):
    name: str = "web_search"
    description: str = "访问互联网获取实时信息，适用于需要最新数据的查询"
    args_schema: Type[BaseModel] = WebSearchInput

    def _run(self, query: str) -> str:
        """执行网络搜索并返回结果"""
        search = DuckDuckGoSearchAPIWrapper()
        return search.run(query)

# 初始化Ollama模型（使用新的集成方式）
ollama_llm = ChatOllama(
    model="qwen2.5:7b",
    temperature=0.7,
    base_url='http://localhost:11434'
)

# 创建工具集
tools = [WebSearchTool()]

# 初始化智能体
agent = initialize_agent(
    tools,
    ollama_llm,
    agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,  # 更适合结构化输入
    verbose=True,
    handle_parsing_errors=True
)

# 示例使用
if __name__ == "__main__":
    query = "2024年巴黎奥运会新增了哪些比赛项目？"
    response = agent.run(query)
    print("最终回答:", response)